from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# 1. 初始化 Ollama 模型连接
llm = OllamaLLM(
    model="deepseek-r1:14b",  # 与 ollama list 中的模型名一致
    base_url="http://localhost:11434",  # Ollama 服务地址
    temperature=0.7,          # 控制随机性 (0-1)
    num_predict=1024,         # 最大生成长度
    top_k=50,                 # 采样参数
    top_p=0.9,                # 采样参数
    stop=["<|EOT|>"],         # DeepSeek 的停止标记
    num_gpu=99,               # 使用全部 GPU 层加速
    num_thread=8,             # CPU 线程数 (4090 可设 8-12)
)

output_parser = StrOutputParser()

# # .from_template 方法创建模板
# prompt = ChatPromptTemplate.from_template("请根据下面的主题，写一篇小红书文章，主题:{topic}！,注意文章要包含小红书标题，和小红书正文部分，语气要亲切，不要太官方正式。")

# chain = prompt | llm | output_parser

# # 2. 执行调用
# result = chain.invoke({"topic": "秀叶牌减肥茶"})
# print(result)


#.from_messages 方法创建模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业的小红书作者，帮助用户写小红书文案，注意文章要包含小红书标题，和小红书正文部分，语气要亲切，不要太官方正式。"),
    ("human", "请根据下面的主题，写一篇小红书文章，主题:{topic}！"),
])

chain = prompt | llm | output_parser

# 2. 执行调用
result = chain.invoke({"topic": "秀叶牌重庆火锅底料"})
print(result)